"""
构建聊天机器人
建立聊天类
1.chatbot类
1)初始大模型和图， _init
2)创建聊天节点 调用大模型 返回结果
3)初始graph  _init  初始化图使用MessagesState的消息 ，保存消息使用RedisSaver
4）运行graph __call
main函数
"""

from model_utils import getLLM
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.redis import RedisSaver
from langgraph.graph import StateGraph,START,END,MessagesState
import uuid
import logging
logging.getLogger("httpx").setLevel(logging.WARNING)

class Chatbot:
    def __init__(self):
        self.memory = RedisSaver("redis://localhost:6379/0")
        self.llm = getLLM()
        self.graph = self.init_graph()

    def __call__(self):
        conf = {"thread_id": "rr123456"}
        while True:
            human = input("我：")
            if human == "bye":
                print("AI:bye")
                break
            r = self.graph.invoke({"messages": [human]}, config=conf)
            print("AI:", r["messages"][-1].content)

    def init_graph(self):
        builder = StateGraph(MessagesState)
        builder.add_node("chat_node", self.chat_node)
        builder.add_edge(START, "chat_node")
        builder.add_edge("chat_node", END)
        return builder.compile(checkpointer=self.memory)

    def chat_node(self, state):
        r = self.llm.invoke(state["messages"])
        return {"messages": [r]}

if __name__ == '__main__':
    chat = Chatbot()
    chat()